import pyfits, pywcs
import matplotlib.pyplot as plt
import pywcsgrid2
from pywcsgrid2.axes_grid import AxesGrid
from pywcsgrid2.axes_wcs import AxesWcs, GridHelperWcs
import re

wcs_key_pattern = re.compile(r'^(NAXIS|CD|CDELT|CRPIX|CRVAL|CTYPE|CROTA|LONGPOLE|LATPOLE|PV|DISTORT|OBJECT|BUNIT|EPOCH|EQUINOX|LTV|LTM|DTV|DTM)')

def mywcs(h):
    l = [card for card in h.ascardlist() if wcs_key_pattern.match(card.key)]
    return pywcs.WCS(pyfits.Header(l))


if 1:
    F = plt.figure(1, (9, 3.5))

    fnames = ["ICAr2_bin0.5_00.raw.fits",
              "ICAr2_bin0.5_01.raw.fits",
              "ICAr2_bin0.5_02.raw.fits"]
    f = pyfits.open(fnames[0])
    d, h = f[0].data, f[0].header
    #wcs = pywcs.WCS(h)
    wcs = mywcs(h)

    import matplotlib.colorbar

    grid_helper = GridHelperWcs(wcs)
    grid = AxesGrid(F, 111, # similar to subplot(111)
                    nrows_ncols = (2, 2),
                    direction="row",
                    axes_pad = 0.05,
                    add_all=True,
                    label_mode = "L",
                    share_all=True,
                    aspect=True,
                    axes_class = AxesWcs,
                    axes_class_args = dict(grid_helper=grid_helper),
                    )

    from scipy.ndimage import gaussian_filter
    rgb = []
    for i, fname in enumerate(fnames):
        f = pyfits.open(fname)
        d, h = f[0].data, f[0].header
        g = gaussian_filter(d.astype("d"), 1.)
        grid[i].imshow(g, origin="lower")
        #cmap=matplotlib.colorbar.cm.gray_r)
        mymax, mymin = g.max(), g.min()/1.5
        g[g<mymin] = mymin
        g[g>mymax] = mymax
        rgb.append((g-mymin)/(mymax-mymin))


    grid[3].imshow(np.transpose(rgb,[1, 2, 0]), origin="lower")

    grid[1].axislines["left"].major_ticklabels.set_visible(False)
    grid[3].axislines["left"].major_ticklabels.set_visible(False)
    grid[0].axislines["bottom"].major_ticklabels.set_visible(False)
    grid[1].axislines["bottom"].major_ticklabels.set_visible(False)

    grid[2].axislines["bottom"].label.set_text(r"$\alpha_{2000}$")
    grid[2].axislines["left"].label.set_text(r"$\delta_{2000}$")

    grid[0].set_xlim(128, 347)
    grid[0].set_ylim(114, 340)

    hp = grid[0].get_grid_helper()
    hp.update_wcsgrid_params(label_density=(4,4))


